package cn.xxliang.aiproject.config;

import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.OllamaEmbeddingClient;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;



/**
 * @author xxliang
 * @date 2025/7/4  21:11
 * @description
 */
@Configuration
@ComponentScan("cn.xxliang.aiproject")
public class OllamaConfig {

    @Bean
    public OllamaApi ollamaApi(@Value("${spring.ai.ollama.base-url}") String baseUrl) {
        OllamaApi ollamaApi = new OllamaApi(baseUrl);
        return ollamaApi;
    }

    @Bean
    public OpenAiApi openAiApi(@Value("${spring.ai.openai.base-url}") String baseUrl, @Value("${spring.ai.openai.api-key}") String apiKey) {
        OpenAiApi openAiApi = new OpenAiApi(baseUrl, apiKey);
        return openAiApi;
    }



    @Bean
    public OllamaChatClient ollamaChatClient(OllamaApi ollamaApi) {
        OllamaChatClient ollamaChatClient = new OllamaChatClient(ollamaApi);
        return ollamaChatClient;
    }

    @Bean
    public OpenAiChatClient openAiChatClient(OpenAiApi openAiApi) {
        OpenAiChatClient openAiChatClient = new OpenAiChatClient(openAiApi);
        return openAiChatClient;
    }
    /**
     * 创建分词器
     * @return
     */
    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    /**
     * 创建向量数据库
     * @return
     */
    @Bean
    public SimpleVectorStore simpleVectorStore(@Value("${spring.ai.rag.embed}")String model, OllamaApi ollamaApi,OpenAiApi openAiApi){

        //#nomic-embed-text、text-embedding-ada-002
        if("nomic-embed-text".equals( model)){
            OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
            embeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
            return new SimpleVectorStore(embeddingClient);
        }
        else{
            OpenAiEmbeddingClient embeddingClient = new OpenAiEmbeddingClient(openAiApi);
            return new SimpleVectorStore(embeddingClient);
        }
    }

    /**
     * 创建pgVectorStore,不同的模型尽量不要使用相同的向量数据库；
     * @param model
     * @param ollamaApi
     * @param openAiApi
     * @param jdbcTemplate
     * @return
     */
    @Bean
    public PgVectorStore pgVectorStore(@Value("${spring.ai.rag.embed}") String model, OllamaApi ollamaApi, OpenAiApi openAiApi, JdbcTemplate jdbcTemplate) {
        if ("nomic-embed-text".equalsIgnoreCase(model)) {
            OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
            embeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
            return new PgVectorStore(jdbcTemplate, embeddingClient);
        } else {
            OpenAiEmbeddingClient embeddingClient = new OpenAiEmbeddingClient(openAiApi);
            return new PgVectorStore(jdbcTemplate, embeddingClient);
        }
    }


//    @Bean
//    public SimpleVectorStore simpleVectorStore(OllamaApi ollamaApi){
//        OllamaEmbeddingClient ollamaEmbeddingClient = new OllamaEmbeddingClient(ollamaApi);
//        // 设置 ollama的向量模型
//        ollamaEmbeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
//        return new SimpleVectorStore(ollamaEmbeddingClient);
//    }
//
//    /**
//     * 创建pgVectorStore
//     * @param ollamaApi
//     * @param jdbcTemplate
//     * @return
//     */
//    @Bean
//    public PgVectorStore pgVectorStore(OllamaApi ollamaApi, JdbcTemplate jdbcTemplate){
//        OllamaEmbeddingClient ollamaEmbeddingClient = new OllamaEmbeddingClient(ollamaApi);
//        ollamaEmbeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
//        // 创建pgVectorStore
//        return new PgVectorStore(jdbcTemplate, ollamaEmbeddingClient);
//    }



}
